import onnxruntime
from PIL import  Image
import numpy as np
from .keys import alphabetChinese as alphabet


from .util import strLabelConverter, resizeNormalize

converter = strLabelConverter(''.join(alphabet))

def softmax(x):
    x_row_max = x.max(axis=-1)
    x_row_max = x_row_max.reshape(list(x.shape)[:-1]+[1])
    x = x - x_row_max
    x_exp = np.exp(x)
    x_exp_row_sum = x_exp.sum(axis=-1).reshape(list(x.shape)[:-1]+[1])
    softmax = x_exp / x_exp_row_sum
    return softmax


class CRNNHandle:
    def __init__(self, model_path):

        self.sess = onnxruntime.InferenceSession(model_path)

    def predict(self, image):
        """
        预测
        """
        scale = image.size[1] * 1.0 / 32
        w = image.size[0] / scale
        w = int(w)
        transformer = resizeNormalize((w, 32))

        image = transformer(image)

        image = image.transpose(2, 0, 1)
        transformed_image = np.expand_dims(image, axis=0)

        preds = self.sess.run(["out"], {"input": transformed_image.astype(np.float32)})
        preds = preds[0]


        length  = preds.shape[0]
        preds = preds.reshape(length,-1)

        preds = np.argmax(preds,axis=1)

        preds = preds.reshape(-1)


        sim_pred = converter.decode(preds, length, raw=False)

        return sim_pred


    def predict_rbg(self, im, whitelist):
        """
        预测
        """
        scale = im.size[1] * 1.0 / 32
        w = im.size[0] / scale
        w = int(w)

        img = im.resize((w, 32), Image.BILINEAR)
        img = np.array(img, dtype=np.float32)
        img -= 127.5
        img /= 127.5
        image = img.transpose(2, 0, 1)
        transformed_image = np.expand_dims(image, axis=0)

        preds = self.sess.run(["out"], {"input": transformed_image.astype(np.float32)})
        preds = preds[0]

        length  = preds.shape[0]
        preds = preds.reshape(length,-1)

        if whitelist:
            sim_pred = []
            for i in range(length):
                tmp = np.argmax(preds[i])
                if (i<length-1 and tmp != np.argmax(preds[i+1])) or i==length-1:
                    while True:
                        if tmp != 0:
                            str = converter.decode_one(tmp)
                            if str in whitelist:
                                sim_pred.append(str)
                                break
                            else:
                                preds[i][tmp] = -1000
                        else:
                            break
                        tmp = np.argmax(preds[i])
            return ''.join(sim_pred)
        else:
            preds = np.argmax(preds,axis=1)
            preds = preds.reshape(-1)
            sim_pred = converter.decode(preds, length, raw=False)
            return sim_pred

if __name__ == "__main__":
    im = Image.open("471594277244_.pic.jpg")
    crnn_handle = CRNNHandle(model_path="../../data/models/crnn_lite_lstm_bk.onnx")
    print(crnn_handle.predict(im))
